{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tree-Player Dungeons & Dragons with GigaChat and GPT\n",
    "\n",
    "In this notebook, we show how we can use concepts from [CAMEL](https://www.camel-ai.org/) to simulate a role-playing game with a protagonist and a dungeon master. To simulate this game, we create an `DialogueSimulator` class that coordinates the dialogue between the two agents."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import LangChain related modules "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Callable, Dict, List\n",
    "\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chat_models.gigachat import GigaChat\n",
    "from langchain.schema import (\n",
    "    HumanMessage,\n",
    "    SystemMessage,\n",
    ")\n",
    "\n",
    "# По возможности используйте самые большие модели из доступных.\n",
    "giga = GigaChat(profanity=False, temperature=0.2, credentials=...)\n",
    "llm = ChatOpenAI(\n",
    "    model=\"gpt-3.5-turbo\", temperature=0.4, openai_api_key=\"<openai_api_key>\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `DialogueAgent` class\n",
    "The `DialogueAgent` class is a simple wrapper around the `ChatOpenAI` model that stores the message history from the `dialogue_agent`'s point of view by simply concatenating the messages as strings.\n",
    "\n",
    "It exposes two methods: \n",
    "- `send()`: applies the chatmodel to the message history and returns the message string\n",
    "- `receive(name, message)`: adds the `message` spoken by `name` to message history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DialogueAgent:\n",
    "    def __init__(\n",
    "        self,\n",
    "        name: str,\n",
    "        system_message: SystemMessage,\n",
    "        model: llm,\n",
    "    ) -> None:\n",
    "        self.name = name\n",
    "        self.system_message = system_message\n",
    "        self.model = model\n",
    "        self.prefix = f\"{self.name}:\"\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.message_history = [\"Вот разговор:\"]\n",
    "\n",
    "    def send(self) -> str:\n",
    "        \"\"\"\n",
    "        Applies the chatmodel to the message history\n",
    "        and returns the message string\n",
    "        \"\"\"\n",
    "        message = self.model(\n",
    "            [\n",
    "                self.system_message,\n",
    "                HumanMessage(content=\"\\n\".join(self.message_history + [self.prefix])),\n",
    "            ]\n",
    "        )\n",
    "        return message.content\n",
    "\n",
    "    def receive(self, name: str, message: str) -> None:\n",
    "        \"\"\"\n",
    "        Concatenates {message} spoken by {name} into message history\n",
    "        \"\"\"\n",
    "        self.message_history.append(f\"{name}: {message}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `DialogueSimulator` class\n",
    "The `DialogueSimulator` class takes a list of agents. At each step, it performs the following:\n",
    "1. Select the next speaker\n",
    "2. Calls the next speaker to send a message \n",
    "3. Broadcasts the message to all other agents\n",
    "4. Update the step counter.\n",
    "The selection of the next speaker can be implemented as any function, but in this case we simply loop through the agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DialogueSimulator:\n",
    "    def __init__(\n",
    "        self,\n",
    "        agents: List[DialogueAgent],\n",
    "        selection_function: Callable[[int, List[DialogueAgent]], int],\n",
    "    ) -> None:\n",
    "        self.agents = agents\n",
    "        self._step = 0\n",
    "        self.select_next_speaker = selection_function\n",
    "\n",
    "    def reset(self):\n",
    "        for agent in self.agents:\n",
    "            agent.reset()\n",
    "\n",
    "    def inject(self, name: str, message: str):\n",
    "        \"\"\"\n",
    "        Initiates the conversation with a {message} from {name}\n",
    "        \"\"\"\n",
    "        for agent in self.agents:\n",
    "            agent.receive(name, message)\n",
    "\n",
    "        # increment time\n",
    "        self._step += 1\n",
    "\n",
    "    def step(self) -> tuple[str, str]:\n",
    "        # 1. choose the next speaker\n",
    "        speaker_idx = self.select_next_speaker(self._step, self.agents)\n",
    "        speaker = self.agents[speaker_idx]\n",
    "\n",
    "        # 2. next speaker sends message\n",
    "        message = speaker.send()\n",
    "\n",
    "        # 3. everyone receives message\n",
    "        for receiver in self.agents:\n",
    "            receiver.receive(speaker.name, message)\n",
    "\n",
    "        # 4. increment time\n",
    "        self._step += 1\n",
    "\n",
    "        return speaker.name, message"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define roles and quest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "protagonist_name = \"Гарри Поттер\"\n",
    "storyteller_name = \"Старый маг\"\n",
    "thirdparty_name = \"Призрак-шутник\"\n",
    "quest = \"Найти все 50 крестражей Кощея Бессмертного\"\n",
    "word_limit = 50  # word limit for task brainstorming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ask an LLM to add detail to the game description"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_description = f\"\"\"Мы будем играть в Dungeons & Dragons. Задание: {quest}.\n",
    "        Главный игрок - {protagonist_name}.\n",
    "        Второй игрок - {thirdparty_name} комментирует происходящее шутками, но не участвует в сюжете.\n",
    "        Ведущий - расказчик, {storyteller_name}.\"\"\"\n",
    "\n",
    "player_descriptor_system_message = SystemMessage(\n",
    "    content=\"Действие происходит во вселенной Гарри Поттера, куда проник герой русских народных сказок, Кощей Бессмертный.\"\n",
    ")\n",
    "\n",
    "protagonist_specifier_prompt = [\n",
    "    player_descriptor_system_message,\n",
    "    HumanMessage(\n",
    "        content=f\"\"\"{game_description}\n",
    "        Пожалуйста предложи креативное описание главного героя по имени {protagonist_name}, уложись в {word_limit} слов или меньше. \n",
    "        Говори напрямую с {protagonist_name}.\n",
    "        Больше ничего не добавляй\"\"\"\n",
    "    ),\n",
    "]\n",
    "\n",
    "protagonist_description = llm(protagonist_specifier_prompt).content\n",
    "\n",
    "storyteller_specifier_prompt = [\n",
    "    player_descriptor_system_message,\n",
    "    HumanMessage(\n",
    "        content=f\"\"\"{game_description}\n",
    "        Пожалуйста предложи креативное описание рассказчика по имени {storyteller_name}, уложись в {word_limit} слов или меньше. \n",
    "        Говори напрямую с {storyteller_name}.\n",
    "        Больше ничего не добавляй.\"\"\"\n",
    "    ),\n",
    "]\n",
    "\n",
    "storyteller_description = llm(storyteller_specifier_prompt).content\n",
    "\n",
    "thirdparty_specifier_prompt = [\n",
    "    player_descriptor_system_message,\n",
    "    HumanMessage(\n",
    "        content=f\"\"\"{game_description}\n",
    "        Пожалуйста предложи креативное описание комментатора по имени {thirdparty_name}, уложись в {word_limit} слов или меньше. \n",
    "        Говори напрямую с {thirdparty_name}.\n",
    "        Больше ничего не добавляй\"\"\"\n",
    "    ),\n",
    "]\n",
    "\n",
    "thirdparty_description = llm(thirdparty_specifier_prompt).content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[33mВедущий:\u001b[0m\n",
      "Старый маг, седовласый и мудрый, в глазах его горит огонь древних знаний. В его голосе звучит магическая мелодия, словно он сам является живым заклинанием. Он владеет силой слова и умеет увлечь слушателей в мир сказок и приключений, где реальность и волшебство переплетаются в неразрывную сеть.\n",
      "\n",
      "\u001b[32mГлавный игрок:\u001b[0m\n",
      "Гарри Поттер, юный волшебник с метлой в руке и огонь в сердце. В его глазах светится решимость и непокорность, а волосы, словно молнии, олицетворяют его силу и страсть к приключениям. Он – символ надежды и смелости, готовый сразиться с любым злом, чтобы защитить своих друзей и мир магии.\n",
      "\n",
      "\u001b[31mКомментатор:\u001b[0m\n",
      "Призрак-шутник, сияющий смехом и безумными глазами, скитается по миру, окутывая все вокруг своими шутками и розыгрышами. Он словно живая комедия, которая никогда не заканчивается. Его голос звучит игриво и проникает в самые глубины сознания, заставляя улыбаться и забывать о реальности.\n"
     ]
    }
   ],
   "source": [
    "# Yellow color text\n",
    "print(\"\\n\\033[33mВедущий:\\033[0m\")\n",
    "print(storyteller_description)\n",
    "\n",
    "# Green color text\n",
    "print(\"\\n\\033[32mГлавный игрок:\\033[0m\")\n",
    "print(protagonist_description)\n",
    "\n",
    "# Red text\n",
    "print(\"\\n\\033[31mКомментатор:\\033[0m\")\n",
    "print(thirdparty_description)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Protagonist and dungeon master system messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "protagonist_system_message = SystemMessage(\n",
    "    content=(\n",
    "        f\"\"\"{game_description}\n",
    "Никогда не забывай, что ты главный игрок - {protagonist_name}, а я - рассказчик, {storyteller_name}. \n",
    "Вот описание твоего персонажа: {protagonist_description}.\n",
    "Ты предлагаешь действия, которые планируешь предпринять, и я объясню, что произойдет, когда ты предпримешь эти действия.\n",
    "Говори в первом лице от имени персонажа {protagonist_name}.\n",
    "Для описания движений собственного тела заключите описание в «*».\n",
    "Не меняй роли!\n",
    "Не говори с точки зрения персонажа {storyteller_name}.\n",
    "Больше ничего не добавляй.\n",
    "Запомни, что ты главный герой - {protagonist_name}.\n",
    "Прекращай говорить, когда тебе кажется, что ты закончил мысль.\n",
    "Отвечай коротко, одной строкой, уложись в {word_limit} слов или меньше.\n",
    "Не отвечай одно и то же! Развивай историю, совершай новые действия. Читателю должно быть интересно. Придумывай новые трудности для персонажа {protagonist_name} и поменьше разговаривай с ним.\n",
    "\"\"\"\n",
    "    )\n",
    ")\n",
    "\n",
    "storyteller_system_message = SystemMessage(\n",
    "    content=(\n",
    "        f\"\"\"{game_description}\n",
    "Никогда не забывай, что ты рассказчик - {storyteller_name}, а я главный герой - {protagonist_name}. \n",
    "Вот описание твоего персонажа: {storyteller_description}.\n",
    "Ты предлагаешь действия, которые планируешь предпринять, и я объясню, что произойдет, когда ты предпримешь эти действия.\n",
    "Говори в первом лице от имени персонажа {storyteller_name}.\n",
    "Для описания движений собственного тела заключите описание в «*».\n",
    "Не меняй роли!\n",
    "Не говори с точки зрения персонажа {protagonist_name}.\n",
    "Больше ничего не добавляй.\n",
    "Запомни, что ты рассказчик - {storyteller_name}.\n",
    "Прекращай говорить, когда тебе кажется, что ты закончил мысль.\n",
    "Отвечай коротко, одной строкой, уложись в {word_limit} слов или меньше.\n",
    "Не отвечай одно и то же! Развивай историю, совершай новые активные действия. Меньше говори и больше действуй, чтобы сюжет развивался. Читателю должно быть интересно.\n",
    "\"\"\"\n",
    "    )\n",
    ")\n",
    "\n",
    "thirdparty_system_message = SystemMessage(\n",
    "    content=(\n",
    "        f\"\"\"{game_description}\n",
    "Ты - {thirdparty_name}.\n",
    "Вот описание твоего персонажа - {thirdparty_description}\n",
    "Напиши шутку про происходящее. Пиши строго не более 5 слов и ничего больше.\n",
    "\"\"\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use an LLM to create an elaborate quest description"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Original quest:\n",
      "Найти все 50 крестражей Кощея Бессмертного\n",
      "\n",
      "Detailed quest:\n",
      "Добро пожаловать, молодой Гарри Поттер! Великая задача ожидает тебя в мире Dungeons & Dragons. Ты - наш главный герой, судьба Кощея Бессмертного лежит в твоих руках. Старый маг приветствует тебя и желает удачи в поиске 50 крестражей. Пусть твоя магия будет сильной и твои приключения запоминающимися!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "quest_specifier_prompt = [\n",
    "    SystemMessage(\n",
    "        content=f\"Напиши стартовую реплику для истории от имени {storyteller_name}, который обращается к {protagonist_name}\"\n",
    "    ),\n",
    "    HumanMessage(\n",
    "        content=f\"\"\"{game_description}\n",
    "        \n",
    "        Ты - рассказчик, {storyteller_name}.\n",
    "        Пожалуйста, напиши вводную фразу для начала истории. Будь изобретатен и креативен.\n",
    "        Пожалуйста, уложись в {word_limit} слов или меньше\n",
    "        Обращайся непосредственно к персонажу {protagonist_name}.\"\"\"\n",
    "    ),\n",
    "]\n",
    "specified_quest = llm(quest_specifier_prompt).content\n",
    "\n",
    "print(f\"\\nOriginal quest:\\n{quest}\\n\")\n",
    "print(f\"Detailed quest:\\n{specified_quest}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "protagonist = DialogueAgent(\n",
    "    name=protagonist_name,\n",
    "    system_message=protagonist_system_message,\n",
    "    model=llm,\n",
    ")\n",
    "storyteller = DialogueAgent(\n",
    "    name=storyteller_name,\n",
    "    system_message=storyteller_system_message,\n",
    "    model=llm,\n",
    ")\n",
    "thirdparty = DialogueAgent(\n",
    "    name=thirdparty_name,\n",
    "    system_message=thirdparty_system_message,\n",
    "    model=giga,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_next_speaker(step: int, agents: List[DialogueAgent]) -> int:\n",
    "    idx = step % len(agents)\n",
    "    return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33m(Старый маг): Добро пожаловать, молодой Гарри Поттер! Великая задача ожидает тебя в мире Dungeons & Dragons. Ты - наш главный герой, судьба Кощея Бессмертного лежит в твоих руках. Старый маг приветствует тебя и желает удачи в поиске 50 крестражей. Пусть твоя магия будет сильной и твои приключения запоминающимися!\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[32m(Гарри Поттер): Спасибо, Старый маг! Я готов принять вызов и отправиться в поиски крестражей Кощея Бессмертного. Я буду использовать свою метлу и волшебные способности, чтобы преодолеть все трудности и вернуть миру магии его равновесие. Приключение начинается!\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[31m(Призрак-шутник): Призрак-шутник: Ха-ха-ха, я знал, что ты справишься, Гарри Поттер! Удачи в поисках крестражей!\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[33m(Старый маг): Гарри Поттер, ты отправляешься в первый город, где, по легенде, находится один из крестражей Кощея Бессмертного. Ты видишь перед собой старинные улочки, торговые лавки и толпу народа. Куда ты направишься?\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[32m(Гарри Поттер): Я направлюсь к местной таверне, чтобы пообщаться с местными жителями и узнать, есть ли у них какая-либо информация о местонахождении крестража Кощея Бессмертного.\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[31m(Призрак-шутник): Призрак-шутник: Ха-ха-ха, Гарри Поттер, они расскажут тебе о своих тайнах и секретах, а потом продадут тебе за бесценок какой-нибудь ненужный артефакт!\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[33m(Старый маг): Ты входишь в таверну, где тебя встречает атмосфера шума и смеха. Ты подходишь к барной стойке и заказываешь напиток. Рядом с тобой сидит старый мужчина, который кажется знающим. Что ты сделаешь?\u001b[0m\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "max_iters = 6\n",
    "n = 0\n",
    "\n",
    "simulator = DialogueSimulator(\n",
    "    agents=[storyteller, protagonist, thirdparty],\n",
    "    selection_function=select_next_speaker,\n",
    ")\n",
    "\n",
    "simulator.reset()\n",
    "simulator.inject(storyteller_name, specified_quest)\n",
    "print(f\"\\033[33m({storyteller_name}): {specified_quest}\\033[0m\")\n",
    "print(\"\\n\")\n",
    "\n",
    "while n < max_iters:\n",
    "    name, message = simulator.step()\n",
    "    # Make player green\n",
    "    if name == protagonist_name:\n",
    "        print(f\"\\033[32m({name}): {message}\\033[0m\")\n",
    "    elif name == storyteller_name:  # Yellow\n",
    "        print(f\"\\033[33m({name}): {message}\\033[0m\")\n",
    "    elif name == thirdparty_name:  # Red\n",
    "        print(f\"\\033[31m({name}): {message}\\033[0m\")\n",
    "    print(\"\\n\")\n",
    "    n += 1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
